from __future__ import annotations
from collections import deque
from threading import Thread, Timer, Lock
from abc import ABC, abstractmethod
from concurrent.futures import Executor
from typing import Union, List, Tuple, Callable, Any, Literal, TypeVar, Generic, NamedTuple
from inspect import iscoroutinefunction
import time, traceback
import asyncio

T_in = TypeVar('T_in')
T_out = TypeVar('T_out')



_global_recorder = {}
import os

class BaseStream(ABC, Generic[T_in, T_out]):
  """_summary_
  流模块abc, 继承后需要实现__call__
  分为 协程异步模式 和 计时器线程轮询模式
  协程异步模式消耗少, 需手动设置线程池并根据线程池分配任务, 需要在async中运行
  计时器线程轮询模式消耗多, 每个轮询开启一个线程
  两者都受到GIL影响, cpu密集任务尽量调用外部库, 不要用python原生语言
  尽量使用协程异步模式(对异步协程较为理解的话), 简单任务可以使用计时器线程轮询模式(或者需要debug观察队列情况时)
  基本的BaseStream流程
  ______________________________    
  |  Getter(=Upstream_acquired)|    
  |  Processor(__call__)       |    
  |  Output(data & signal)     |    
  |  Putter(put)               | -> call(on_full)
  |  Deque                     |    
  |__Exporter(acquired)________| -> call(on_empty)
  继承本模块先继承BasicType, 添加类型检查
  再实现本模块__call__如何处理T_in和T_out
  特殊情况则重构 on_empty和on_full 
  """
  
  _is_working = False # 放在头部避免继承类是单例时会刷新
  _logger = None
  # TODO: 将上下游引用换成弱引用, 维护一个上下游引用字典, 通过property暴露上下游属性
  
  def __init__(self, 
               type_in: T_in,
               type_out: T_out,
               *,
               executor: Executor = None,
               order_type: Literal['FIFO', 'LIFO'] = 'FIFO',
               max_queue_size: int = None,
               interval_ms: int = 10, # 设置Timer / 0时不使用（只能用协程async）
               block: bool = True, # 设置get空时是否堵塞/否则向下传None
               up_stream: BaseStream = None, # 接收来源只能是一个
               dn_stream: BaseStream = None, # 接收来源只能是一个(方便链式设置)
               ):
    if interval_ms <= 0:
      raise Exception("Timer interval must be greater than 0")
    
    self._lock = Lock() # 线程模式下需要锁保证线程安全
    self._executor = executor
    self.is_coroutine = iscoroutinefunction(self.__call__)
    
    assert order_type in ['FIFO', 'LIFO']
    self._order_type = order_type
    
    self._type_in = type_in
    self._type_out = type_out
    self._up_stream = up_stream
    self._dn_stream = dn_stream
    
    self._is_block = block
    self._deque = deque(maxlen=max_queue_size) # 核心缓冲区
    
    self._timer_interval = interval_ms / 1000
    self._timer: Timer = None
    
  @property
  def deque(self):
    return self._deque
  
  @deque.setter
  def deque(self, value):
    raise Exception("deque is read-only")
  
  @deque.deleter
  def deque(self):
    raise Exception("deque is read-only")
  
  @property
  def is_working(self):
    return self._is_working
  
  @property
  def checkin(self):
    return (self._type_in, type(None))
  
  @property
  def checkout(self):
    return (self._type_out, type(None))
  
  def clear(self):
    """
    清空队列
    """
    with self._lock:
      self._deque.clear()
    
  def start(self):
    if self._dn_stream is not None:
      self._dn_stream.start()
    self._is_working = True
    if self._up_stream is None:
      self._timer = Timer(self._timer_interval, self.send)
    else:
      self._timer = Timer(self._timer_interval, self.flow)
    setattr(self._timer, 'daemon', True)
    self._timer.start()
    
  def collect_dn_stream_flow(self, 
                             tasks: List[asyncio.Task] = [],
                             run_forever: bool=True):
    if run_forever:
      tasks.append(self.async_loop())
    else:
      tasks.append(self.async_flow())
    if self._dn_stream is not None:
      return self._dn_stream.collect_dn_stream_flow(tasks, run_forever)
    return tasks
  
  async def async_loop(self):
    self._is_block = False # 协程模式强制不可堵塞
    self._is_working = True
    while self._is_working:
      try:
        if self._up_stream is None:
          await self.async_send(None)
        else:
          await self.async_flow()
        await asyncio.sleep(0.001) # 需要触发交出控制权
      except asyncio.CancelledError:
        self._is_working = False
        break
      except Exception as e:
        print(f"Flow error: {self.__class__.__name__} {str(e)}")
        traceback.print_exc()
        raise e
    self._is_working = False
    print(f"Flow exit: {self.__class__.__name__}")
    
  async def async_start(self, 
                        run_forever: bool=True):
    """
        该方法必须仅由最上游流模块调用
    """
    self._is_block = False # 协程模式强制不可堵塞
    self._is_working = True
    tasks = [asyncio.create_task(task) for task in self.collect_dn_stream_flow(run_forever=run_forever)]
    await asyncio.gather(*tasks)
    
  def stop(self):
    """
      如果存在需要释放的句柄, 在这里实现 
    """
    if self._dn_stream is not None:
      self._dn_stream.stop()
    self._is_working = False
    if cancel := getattr(self._timer, 'cancel', None):
      cancel()
    
  @property
  def pause(self):
    """
    暂停上下文
    """
    class PauseContext:
      def __enter__(this):
        if cancel := getattr(self._timer, 'cancel', None):
          cancel()
        
      def __exit__(this, exc_type, exc_value, traceback):
        if not exc_type is None:
          self._is_working = False
          raise exc_value
        if self._is_working and self._timer_interval > 0:
          if self._up_stream is None:
            self._timer = Timer(self._timer_interval, self.send)
          else:
            self._timer = Timer(self._timer_interval, self.flow)
          setattr(self._timer, 'daemon', True)
          self._timer.start()
    return PauseContext()
    
  def hook_to(self,
              module: BaseStream,
              ) -> BaseStream:
    """
    设置上游
    """
    if self._up_stream is not None:
      if self._up_stream._dn_stream == self:
        return
      raise Exception("Only one upstream module is allowed")
    self._up_stream = module
    module.link_to(self)
    return module
    
  def link_to(self,
              module: BaseStream,
              ) -> BaseStream:
    """
    设置下游
    """
    if self._dn_stream is not None:
      if self._dn_stream._up_stream == self:
        return
      raise Exception("Only one downstream module is allowed")
    self._dn_stream = module
    module.hook_to(self)
    return module
  
  async def async_acquired(self
                      ) -> T_out:
    """
    本模块被获取数据时调用
    一般在FLOW中是被下游模块主动调用
    异步应用使用该方法
    """
    loop = asyncio.get_running_loop()
    if len(self._deque) == 0:
      # 处理空队列
      data = await loop.run_in_executor(self._executor, 
                                        self.on_empty, 
                                        self._deque)
    else:
      with self._lock:
        if self._order_type == 'FIFO':
          data = self._deque.popleft()
        else:
          data = self._deque.pop()
    assert isinstance(data, self.checkout)
    return data
  
  def acquired(self
          ) -> T_out:
    """
    本模块被获取数据时调用
    一般在FLOW中是被下游模块主动调用
    同步应用使用该方法
    """
    if len(self._deque) == 0:
      data = self.on_empty(self._deque)
    else:
      with self._lock:
        if self._order_type == 'FIFO':
          data = self._deque.popleft()
        else:
          data = self._deque.pop()
    assert isinstance(data, self.checkout)
    return data
  
  async def async_put(self, 
                      data: T_out,
                      include_list: bool = True,
                      ) :
    checkout = self.checkout + (list, ) if include_list else self.checkout
    assert isinstance(data, checkout)
    if data is None: # None默认不处理
      return
    if isinstance(data, list) and include_list:
      for item in data:
        await self.async_put(item, False) # 谨防无限递归
      return 
    if maxlen:= self._deque.maxlen:
      if len(self._deque) + 1 >= maxlen:
        # 处理满队列
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._executor, 
                                          self.on_full, 
                                          self._deque,
                                          data)
    with self._lock: # TODO: 可以改为异步锁
      self._deque.append(data)
    
  def put(self, 
          data: T_out,
          include_list: bool = True,
          ) :
    checkout = self.checkout + (list, ) if include_list else self.checkout
    assert isinstance(data, checkout)
    if data is None: # None默认不处理
      return
    if isinstance(data, list) and include_list:
      for item in data:
        self.put(item, False) # 谨防无限递归
      return 
    if maxlen:= self._deque.maxlen:
      if len(self._deque) + 1 >= maxlen:
        return self.on_full(self._deque, data)
    with self._lock:
      self._deque.append(data)
      
  async def async_send(self,
                       data: T_in=None,
                       ):
    assert isinstance(data, self.checkin)
    # 最上游模块调用, 默认持续传入空值
    loop = asyncio.get_running_loop()
    if self.is_coroutine:
      output = await self.__call__(data)
    else:
      output = await loop.run_in_executor(self._executor, 
                                          self.__call__,
                                          data)
    await self.async_put(output)
      
  def send(self, 
           data: T_in=None, 
           ): 
    """
    最上游模块需要使用该方法传入数据
    Args:
        data (T_in): 传入数据
    """
    assert isinstance(data, self.checkin)
    # 最上游模块调用, 默认持续传入空值
    with self.pause:
      output = self(data)
      self.put(output)
  
  async def async_flow(self):
    if self._up_stream is None:
      # 无上游时跳过该flow，需要async_send进来
      return 
    else:
      data = await self._up_stream.async_acquired()
    if data is None: # None默认不处理
      return  
    loop = asyncio.get_running_loop()
    if self.is_coroutine:
      output = await self.__call__(data)
    else:
      output = await loop.run_in_executor(self._executor, 
                                          self.__call__,
                                          data)
    await self.async_put(output)
  
  def flow(self):
    with self.pause:
      if self._up_stream is None:
        # 无上游时跳过该flow，需要send进来
        return # 无上游时不运行flow
      else:
        data = self._up_stream.acquired()
      if data is None: # None默认不处理
        return  
      output = self(data)
      self.put(output)
        
  def on_empty(self, 
              deque: deque
              ) -> T_out:
    """
    处理队列为空时的逻辑, 默认逻辑是根据block选择堵塞/无block则丢出None
    最好根据模块逻辑重写该方法
    args:
      deque: 队列
    returns:
      none
    """
    if self._is_block:
      while len(deque) == 0:
        time.sleep(0.01)
      if self._order_type == 'FIFO':
        data = deque.popleft()
      else:
        data = deque.pop()
    else:
      data = None
    return data
          
  def on_full(self, 
              deque: deque,
              data: T_out,
              ):
    """
    处理队列已满时的逻辑, 默认逻辑是根据block选择堵塞/无block则照常执行插入挤出旧数据
    最好根据模块逻辑重写该方法
    args:
      deque: 队列
    returns:
      none
    """
    if self._is_block:
      while len(deque) >= deque.maxlen:
        time.sleep(0.01)
    else:
      if self._order_type == 'FIFO':
        deque.popleft()
      else:
        deque.pop()
      self._deque.append(data)
      
  def get_top_stream(self):
    if self._up_stream is None:
      return self
    else:
      return self._up_stream.get_top_stream()
    
  def get_bottom_stream(self):
    if self._dn_stream is None:
      return self
    else:
      return self._dn_stream.get_bottom_stream()
      
  @abstractmethod
  def __call__(self, 
              data: T_in
              ) -> T_out:
    """
    需要在此处理每次进来的数据如何处理，
    然后放进本模块的队列中
    可以用async def实现, 但是该实现会失去多线程的优势且必须在asyncio模式下执行(仅在需要调用一些异步模块时实现)
    """
    pass
